﻿#include "fixed_mem_pool.hh"

#include "../util/os_util.hh"
#include "../util/singleton.hh"

#include <stdexcept>
#include <string>

//
// 定义宏NO_MEMORY_POOL可以屏蔽内存池，直接使用malloc/free管理内存
//

/**
 * 内存池异常
 */
class MemoryPoolException : public std::exception {
  std::string reason_; ///< 异常原因

public:
  /**
   * 构造
   * \param reason 异常原因
   */
  explicit MemoryPoolException(const char *reason) noexcept {
    reason_ = reason;
  }
  /**
   * 析构
   */
  virtual ~MemoryPoolException() {}
  /**
   * 获取异常原因
   */
  virtual const char *what() const noexcept { return reason_.c_str(); }
};

/**
 * 内存池内分配的内存，头部长度类型.
 */
using prefix_type = int;

/**
 * 内存池配置.
 */
struct MemPoolConfig {
  /**
   * 超过DEFAULT_MIN_COUNT长度限制后不在池内分配.
   * 内存长度将为-1
   */
  constexpr static auto BIG_MEM_MAGIC_NUMBER = prefix_type(-1);
  /**
   * 每个固定内存池内预先分配的内存块个数.
   */
  constexpr static std::size_t DEFAULT_MIN_COUNT = 128;
  /**
   * 每次GC释放的内存块最大数量.
   */
  constexpr static std::size_t RECYCLE_STEP = DEFAULT_MIN_COUNT;
  /**
   * GC检测时间间隔，毫秒.
   */
  constexpr static std::time_t GC_INTVAL = 100;
};

/**
 * 全局的内存池.
 */
static kratos::MemPool mem_pool;

/**
 * 获取全局内存池指针.
 *
 * \return
 */
kratos::MemPool *kratos::get_global_mem_pool() { return &mem_pool; }

/**
 * 内存块信息，位于分配的内存块最前面.
 */
struct MemCookie {
  prefix_type length{0}; ///< 内存块长度
  bool in_use{false};    ///< 是否正在被使用，容错处理
  kratos::FixedMemPool *owner{nullptr}; ///< 内存块所属池
};

kratos::FixedMemPool::FixedMemPool(std::size_t fixed_size) noexcept {
  // 默认配置
  gc_intval_ = MemPoolConfig::GC_INTVAL;
  min_count_ = MemPoolConfig::DEFAULT_MIN_COUNT;
  recycle_step_ = MemPoolConfig::RECYCLE_STEP;

  // 计算内存块实际长度，包含头和实际可用内存
  fixed_size_ = fixed_size + sizeof(MemCookie);
  // 根据配置预先初始化池
  for (std::size_t i = 0; i < min_count_; i++) {
    auto *block = init_one();
    if (block) {
      ptr_list_.push_back(block);
    }
  }
}

kratos::FixedMemPool::~FixedMemPool() noexcept {
  // 销毁池内所有内存块
  for (auto *ptr : ptr_list_) {
    delete[] ptr;
  }
  ptr_list_.clear();
}

auto kratos::FixedMemPool::update(std::time_t now) noexcept -> void {
  if (!last_gc_tick_) {
    last_gc_tick_ = now;
    return;
  }
  if (gc_on_off_) {
    if (now - last_gc_tick_ > gc_intval_) {
      // 启动GC并记录时间戳
      std::scoped_lock guard(lock_);
      step_gc();
      last_gc_tick_ = now;
    }
  }
}

auto kratos::FixedMemPool::init_one() noexcept -> char * {
  char *ptr = nullptr;
  try {
    ptr = new char[fixed_size_];
    if (!ptr) {
      // 内存耗尽
      return nullptr;
    }
  } catch (std::exception &) {
    return nullptr;
  }
  auto *cookie_ptr = reinterpret_cast<MemCookie *>(ptr);
  cookie_ptr->length =
      static_cast<prefix_type>(fixed_size_) - sizeof(MemCookie);
  cookie_ptr->in_use = false;
  cookie_ptr->owner = this;
  return ptr;
}

auto kratos::FixedMemPool::rent() noexcept -> char * {
  char *ptr = nullptr;
  {
    std::scoped_lock guard(lock_);
    if (ptr_list_.empty()) {
      try {
        // 池内耗尽，分配一个新的内存块
        ptr = new char[fixed_size_];
      } catch (std::exception &) {
        // 分配异常
        return nullptr;
      }
      if (!ptr) {
        // 堆内存耗尽
        return nullptr;
      }
    } else {
      // 取出一个可用的内存块
      ptr = ptr_list_.front();
      ptr_list_.pop_front();
    }
    count_in_use_ += 1;
  }
  auto *cookie_ptr = reinterpret_cast<MemCookie *>(ptr);
  cookie_ptr->length =
      static_cast<prefix_type>(fixed_size_) - sizeof(MemCookie);
  cookie_ptr->owner = this;
  cookie_ptr->in_use = true;
  return ptr + sizeof(MemCookie);
}

auto kratos::FixedMemPool::recycle(char *ptr) noexcept(false) -> bool {
  auto *cookie_ptr = reinterpret_cast<MemCookie *>(ptr);
  if (cookie_ptr->owner != this) {
    // 不属于本池
    auto reason = "Recycle, invalid memory ownership[" +
                  std::to_string(cookie_ptr->length) + "]";
    reason += kratos::util::get_current_stack_trace_info();
    throw MemoryPoolException(reason.c_str());
    return false;
  }
  if (cookie_ptr->length != static_cast<prefix_type>(fixed_size_) -
                                static_cast<prefix_type>(sizeof(MemCookie))) {
    // 长度不符
    auto reason = "Recycle, invalid memory length[" +
                  std::to_string(cookie_ptr->length) + "]";
    reason += kratos::util::get_current_stack_trace_info();
    throw MemoryPoolException(reason.c_str());
    return false;
  }
  std::scoped_lock guard(lock_);
  if (!cookie_ptr->in_use) {
    // 防止重复释放
    auto reason =
        "Recycle, double free [" + std::to_string(cookie_ptr->length) + "]\n";
    reason += kratos::util::get_current_stack_trace_info();
    throw MemoryPoolException(reason.c_str());
    return false;
  }
  ptr_list_.push_front(ptr);
  count_in_use_ -= 1;
  cookie_ptr->in_use = false;
  return true;
}

auto kratos::FixedMemPool::count_in_use() const noexcept -> std::size_t {
  return count_in_use_;
}

auto kratos::FixedMemPool::fixed_size() const noexcept -> std::size_t {
  return fixed_size_;
}

auto kratos::FixedMemPool::count_in_pool() noexcept -> std::size_t {
  std::scoped_lock guard(lock_);
  return ptr_list_.size();
}

auto kratos::FixedMemPool::turn_on_off_gc(bool on_off) noexcept -> void {
  gc_on_off_ = on_off;
}

auto kratos::FixedMemPool::step_gc() noexcept -> void {
  std::size_t delete_count = 0;
  while (!ptr_list_.empty() && (ptr_list_.size() >= count_in_use_) &&
         (ptr_list_.size() > min_count_)) {
    if (delete_count >= recycle_step_) {
      break;
    }
    delete_count += 1;
    delete[] ptr_list_.back();
    ptr_list_.pop_back();
  }
}

auto kratos::FixedMemPool::release_all() noexcept -> void {
  std::scoped_lock guard(lock_);
  while (!ptr_list_.empty()) {
    delete[] ptr_list_.front();
    ptr_list_.pop_front();
  }
}

auto kratos::FixedMemPool::set_gc_intval(std::time_t intval) noexcept -> void {
  gc_intval_ = intval;
}

std::size_t next_pow2(std::size_t v) {
  std::size_t p = 1;
  while (p < v) {
    p <<= 1;
  }
  return p;
}

kratos::MemPool::MemPool(std::size_t max_size) noexcept {
#ifdef NO_MEMORY_POOL
  return;
#endif
  gc_intval_ = MemPoolConfig::GC_INTVAL;
  max_size_ = next_pow2(max_size);
  std::size_t index = MIN_SIZE;
  while (index <= max_size_) {
    // 初始化内部多个池
    mem_map_.emplace(index, new FixedMemPool(index));
    index *= 2;
  }
}

kratos::MemPool::~MemPool() noexcept {}

auto kratos::MemPool::update(std::time_t now) noexcept -> void {
#ifdef NO_MEMORY_POOL
  return;
#else
  for (auto &[k, v] : mem_map_) {
    v->update(now);
  }
#endif
}

auto kratos::MemPool::rent(std::size_t size) noexcept(false) -> char * {
#ifdef NO_MEMORY_POOL
  return (char *)malloc((int)size);
#endif
  if (size < MIN_SIZE) {
    size = MIN_SIZE;
  }
  size = next_pow2(size);
  if (size > max_size_) {
    // 超过最大限制，直接从堆上分配
    char *ptr = nullptr;
    try {
      ptr = new char[size + sizeof(MemCookie)];
      if (!ptr) {
        // 内存耗尽
        return nullptr;
      }
    } catch (std::exception &) {
      return nullptr;
    }
    auto *cookie_ptr = reinterpret_cast<MemCookie *>(ptr);
    cookie_ptr->length = MemPoolConfig::BIG_MEM_MAGIC_NUMBER;
    cookie_ptr->in_use = true;
    count_in_use_ += 1;
    return ptr + sizeof(MemCookie);
  } else {
    auto it = mem_map_.find(size);
    if (it == mem_map_.end()) {
      auto reason = "Rent, invalid memory length[" + std::to_string(size) + "]";
      reason += kratos::util::get_current_stack_trace_info();
      throw MemoryPoolException(reason.c_str());
      return nullptr;
    } else {
      auto *rent_ptr = it->second->rent();
      if (rent_ptr) {
        count_in_use_ += 1;
        mem_block_size_in_use_ += it->second->fixed_size();
        return rent_ptr;
      }
    }
  }
  return nullptr;
}

auto kratos::MemPool::recycle(char *ptr) noexcept(false) -> bool {
#ifdef NO_MEMORY_POOL
  free(ptr);
  return true;
#endif
  if (!ptr) {
    return true;
  }
  auto *cookie_ptr = reinterpret_cast<MemCookie *>(ptr - sizeof(MemCookie));
  if (cookie_ptr->length == MemPoolConfig::BIG_MEM_MAGIC_NUMBER) {
    // 超过长度直接释放
    count_in_use_ -= 1;
    delete[](ptr - sizeof(MemCookie));
    return true;
  } else {
    // 回池
    auto it = mem_map_.find(cookie_ptr->length);
    if (it == mem_map_.end()) {
      // 错误
      auto reason = "Recycle, invalid memory length[" +
                    std::to_string(cookie_ptr->length) + "]";
      reason += kratos::util::get_current_stack_trace_info();
      throw MemoryPoolException(reason.c_str());
      return false;
    } else {
      if (it->second->recycle(ptr - sizeof(MemCookie))) {
        mem_block_size_in_use_ -= it->second->fixed_size();
        count_in_use_ -= 1;
        return true;
      }
    }
  }
  return false;
}

auto kratos::MemPool::count_in_use() const noexcept -> std::size_t {
#ifdef NO_MEMORY_POOL
  return 0;
#else
  return count_in_use_;
#endif
}

auto kratos::MemPool::mem_block_size_in_use() const noexcept -> std::size_t {
#ifdef NO_MEMORY_POOL
  return 0;
#else
  return mem_block_size_in_use_;
#endif
}

auto kratos::MemPool::get_fixed_mem_pools() noexcept -> const MemMap & {
  return mem_map_;
}

auto kratos::MemPool::release_all() noexcept -> void {
#ifdef NO_MEMORY_POOL
  return;
#else
  for (auto &[k, v] : mem_map_) {
    v->release_all();
  }
#endif
}

auto kratos::MemPool::get_gc_on_off() const noexcept -> bool {
#ifdef NO_MEMORY_POOL
  return false;
#else
  return gc_on_off_;
#endif
}

auto kratos::MemPool::turn_on_off_gc(bool on_off) noexcept -> void {
#ifdef NO_MEMORY_POOL
  return;
#else
  gc_on_off_ = on_off;
  for (auto &[_, v] : mem_map_) {
    v->turn_on_off_gc(on_off);
  }
#endif
}

auto kratos::MemPool::set_gc_intval(std::time_t intval) noexcept -> void {
#ifdef NO_MEMORY_POOL
  return;
#else
  gc_intval_ = intval;
  for (auto &[_, v] : mem_map_) {
    v->set_gc_intval(intval);
  }
#endif
}

auto kratos::MemPool::get_gc_intval() const noexcept -> std::time_t {
  return gc_intval_;
}
